package other;//import org.dom4j.Document;

import cn.hutool.core.collection.CollectionUtil;
import org.dom4j.Document;
import org.dom4j.Element;
import org.dom4j.Node;
import org.dom4j.io.SAXReader;

import javax.xml.xpath.XPathConstants;
import java.io.*;
import java.util.*;

/**
 * @Auther: Code
 * @Date: 2018/9/9 18:02
 * @Description: 批量转 标记xml为 txt
 */
public class test2 {
    static String xmlDir = "C:\\Users\\chenyb\\Desktop\\89-xml";//文件所在路径，所有文件的根目录，记得修改为你电脑上的文件所在路径
//    static String txtDir = "C:\\Users\\chenyb\\Desktop\\data\\txt";//文件所在路径，所有文件的根目录，记得修改为你电脑上的文件所在路径
    static String txtDir = "C:\\Users\\chenyb\\Desktop\\darknet-master\\data\\90-txt";//文件所在路径，所有文件的根目录，记得修改为你电脑上的文件所在路径

    public static void main(String[] args) throws IOException {

        List<String> allImgName = new ArrayList<>();

        try {
//            recursiveTraversalFolder("xml", xmlDir, allImgName);
            recursiveTraversalFolder("txt", txtDir, allImgName);

            int imgSize = allImgName.size();
            if (allImgName.size() > 0) {
                System.out.println("====== 数据集共有 " + imgSize + " 张图片 ======");
                //打乱 数据集的顺序
//                Collections.shuffle(allImgName);

                StringBuffer trainSb = new StringBuffer();
                StringBuffer validSb = new StringBuffer();

                int splitIndex = (int) (imgSize * 0.7);
                for (int i = 0; i < splitIndex; i++) {
                    trainSb.append(allImgName.get(i));
                }
                for (int i = splitIndex; i < imgSize; i++) {
                    validSb.append(allImgName.get(i));
                }
                //生成 phone_train.txt
                write("C:\\Users\\chenyb\\Desktop\\phone_train.txt", trainSb.toString());
                write("C:\\Users\\chenyb\\Desktop\\phone_valid.txt", validSb.toString());
                System.out.println("====== 生成测试集 " + splitIndex + " 张，测试集 " + (imgSize - splitIndex) + "张 ======");
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * @param key        如果是 “xml” 则 生成 txt 文件，如果是 “txt” 则 生成 phone_train.txt 文件
     * @param path       当前待遍历的 文件路径
     * @param allImgName 生成 phone_train.txt 所需的 allImgName
     */
    public static void recursiveTraversalFolder(String key, String path, List<String> allImgName) throws Exception {
        File folder = new File(path);
        if (folder.exists()) {
            File[] fileArr = folder.listFiles();
//            System.out.println("有 " + fileArr.length + " 个");
            if (null == fileArr || fileArr.length == 0) {
                System.out.println("文件夹是空的!");
                return;
            } else {
                for (File file : fileArr) {
                    if (file.isDirectory()) {//是文件夹，继续递归，如果需要重命名文件夹，这里可以做处理
                        System.out.println("文件夹:" + file.getAbsolutePath() + "，继续递归！");
                        recursiveTraversalFolder(key, file.getAbsolutePath(), allImgName);
                    } else {//是文件
                        // 如果是 xml 则 生成 txt 文件
                        if ("xml".equals(key)) {
                            xml2txt(file, false);
                        }
                        // 如果是 txt 则 生成 phone_train.txt 文件
                        else if ("txt".equals(key)) {
                            getAllImgName(file, allImgName);
                        } else {
                            System.out.println("参数异常！");
                            return;
                        }
                    }
                }
            }
        } else {
            System.out.println("文件不存在!");
        }
    }

    //根据 txt 文件，获取 所有待 遍历的 imgName
    public static void getAllImgName(File file, List<String> allImgName) {
        allImgName.add("data/phone/" + file.getName().replaceFirst("txt", "jpg") + "\n");
    }

    /**
     * 将labelimg标记的xml 提取内容 并转成 Yolov4 的 txt
     *
     * @return 生成的所有txt文件对应的图片 列表
     */
    public static void xml2txt(File file, boolean generateTxt) throws Exception {
        StringBuffer result = new StringBuffer();

        SAXReader reader = new SAXReader();
        Document doc = reader.read(new FileInputStream(file.getPath()));
        Element root = doc.getRootElement();


        Double width = Double.parseDouble(doc.selectSingleNode("//annotation/size/width").getText());
        Double height = Double.parseDouble(doc.selectSingleNode("//annotation/size/height").getText());

        List<Element> list = root.elements("object");
        if (list == null || list.size() <= 0) {
            System.out.println(file.getName() + " 没有检测出 手机！");
            return;
        }
        for (Element node : list) {
            Double xmin = Double.parseDouble(node.element("bndbox").elementText("xmin"));
            Double xmax = Double.parseDouble(node.element("bndbox").elementText("xmax"));
            Double ymin = Double.parseDouble(node.element("bndbox").elementText("ymin"));
            Double ymax = Double.parseDouble(node.element("bndbox").elementText("ymax"));
            double x_center = (xmax + xmin) / 2 / width;
            double y_center = (ymax + ymin) / 2 / height;
            double x_width = (xmax - xmin) / width;
            double y_height = (ymax - ymin) / height;
            String str = String.format("0 %.6f %.6f %.6f %.6f \n", x_center, y_center, x_width, y_height);
            result.append(str);
        }
        String newFileName = file.getName().replaceFirst("xml", "txt");
        if (generateTxt) {
            write(txtDir + "\\" + newFileName, result.toString());
        }
    }

    /**
     * 保存文件
     *
     * @param filepath 待保存的文件路径
     * @param content  待保存的文件内容
     */
    public static void write(String filepath, String content) throws IOException {
        //将写入转化为流的形式
        BufferedWriter bw = new BufferedWriter(new FileWriter(filepath));
        bw.write(content);
        //关闭流
        bw.close();
        System.out.println(filepath + " 完成！");
    }


}